// Copyright (c) 2020 Presto Labs Pte. Ltd.
// Author: taekwon

#pragma once

#include <chrono>
#include <queue>
#include "coin2/base/time.h"

class ITransactionLimiter {
 public:
  virtual bool CheckOnlyWithTimestamp(int64_t timestamp) const = 0;
  virtual bool CheckAndSetWithTimestamp(int64_t timestamp) = 0;
  virtual bool CheckOnly() const = 0;
  virtual bool CheckAndSet() = 0;

  virtual ~ITransactionLimiter() = default;
};

class RateLimiter : public ITransactionLimiter {
 public:
  RateLimiter(int rate_cnt, double per_sec, int64_t base_timestamp) :
      rate_cnt_(rate_cnt),
      per_sec_(per_sec),
      allowance_(rate_cnt),
      last_check_timestamp_(base_timestamp) {
    }

  RateLimiter(int rate_cnt, double per_sec) :
      rate_cnt_(rate_cnt),
      per_sec_(per_sec),
      allowance_(rate_cnt) {
    last_check_timestamp_ = GetCurrentTimestamp();
  }

  bool CheckAndSetWithTimestamp(int64_t current_timestamp) {
    double time_passed_sec = static_cast<double>(current_timestamp - last_check_timestamp_) / 1e9;
    if (time_passed_sec <= 0.0) {
      return false;
    }
    last_check_timestamp_ = current_timestamp;
    allowance_ += time_passed_sec * (rate_cnt_ / per_sec_);
    if (allowance_ > rate_cnt_) {
      allowance_ = rate_cnt_;
    }
    if (allowance_ < rate_discount_) {
      return false;
    }
    allowance_ -= rate_discount_;
    return true;
  }

  bool CheckAndSet() {
    return CheckAndSetWithTimestamp(GetCurrentTimestamp());
  }

  bool CheckOnlyWithTimestamp(int64_t current_timestamp) const {
    double time_passed_sec = static_cast<double>(current_timestamp - last_check_timestamp_) / 1e9;
    if (time_passed_sec <= 0.0) {
      return false;
    }
    double allowance = allowance_ + time_passed_sec * (rate_cnt_ / per_sec_);
    if (allowance > rate_cnt_) {
      allowance = rate_cnt_;
    }
    if (allowance < rate_discount_) {
      return false;
    }
    return true;
  }

  bool CheckOnly() const {
    return CheckOnlyWithTimestamp(GetCurrentTimestamp());
  }

 private:
  const int rate_cnt_;
  const double per_sec_;
  const double rate_discount_ = 1.0;
  double allowance_;
  int64_t last_check_timestamp_;
};

class WindowLimiter : public ITransactionLimiter  {
 public:
  WindowLimiter(int rate_cnt, int64_t window_size_ns) :
      rate_cnt_(rate_cnt),
      window_size_ns_(window_size_ns) {
  }

  bool CheckOnly() const {
    return trans_timestamps_.size() < static_cast<size_t>(rate_cnt_);
  }
  bool CheckOnlyWithTimestamp(int64_t timestamp) const {
    return CheckOnly();
  }
  bool CheckAndSetWithTimestamp(int64_t timestamp) {
    RemoveOldTps(timestamp);
    if (CheckOnly()) {
      AppendTps(timestamp);
      return true;
    }
    return false;
  }
  bool CheckAndSet() {
    return CheckAndSetWithTimestamp(GetCurrentTimestamp());
  }

 private:
  void AppendTps(int64_t timestamp) {
    trans_timestamps_.push(timestamp);
  }
  void RemoveOldTps(int64_t timestamp) {
    int64_t max_timestamp = timestamp - window_size_ns_;
    while (!trans_timestamps_.empty() && trans_timestamps_.front() <= max_timestamp) {
      trans_timestamps_.pop();
    }
  }

  const int rate_cnt_;
  const int64_t window_size_ns_;
  std::queue<int64_t> trans_timestamps_;
};

// not be derived from ITransactionLimiter due to using qty.
class QtyLimiter {
 public:
  QtyLimiter(int rate_cnt, double per_sec, int64_t base_timestamp) :
      rate_cnt_(rate_cnt),
      per_sec_(per_sec),
      allowance_(rate_cnt),
      last_check_timestamp_(base_timestamp) {
    }

  QtyLimiter(int rate_cnt, double per_sec) :
      rate_cnt_(rate_cnt),
      per_sec_(per_sec),
      allowance_(rate_cnt) {
    last_check_timestamp_ = GetCurrentTimestamp();
  }

  bool CheckAndSetWithQtyTimestamp(double qty, int64_t current_timestamp) {
    double time_passed_sec = static_cast<double>(current_timestamp - last_check_timestamp_) / 1e9;
    if (time_passed_sec <= 0.0) {
      return false;
    }
    last_check_timestamp_ = current_timestamp;
    allowance_ += time_passed_sec * (rate_cnt_ / per_sec_);
    if (allowance_ > rate_cnt_) {
      allowance_ = rate_cnt_;
    }
    if (allowance_ < qty) {
      return false;
    }
    allowance_ -= qty;
    return true;
  }

  bool CheckAndSetWithQty(double qty) {
    return CheckAndSetWithQtyTimestamp(qty, GetCurrentTimestamp());
  }

  bool CheckOnlyWithQtyTimestamp(double qty, int64_t current_timestamp) const {
    double time_passed_sec = static_cast<double>(current_timestamp - last_check_timestamp_) / 1e9;
    if (time_passed_sec <= 0.0) {
      return false;
    }
    double allowance = allowance_ + time_passed_sec * (rate_cnt_ / per_sec_);
    if (allowance > rate_cnt_) {
      allowance = rate_cnt_;
    }
    if (allowance < qty) {
      return false;
    }
    return true;
  }

  bool CheckOnlyWithQty(double qty) const {
    return CheckOnlyWithQtyTimestamp(qty, GetCurrentTimestamp());
  }

 private:
  const int rate_cnt_;
  const double per_sec_;
  double allowance_;
  int64_t last_check_timestamp_;
};
